package com.docchatlink.iflytek2;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.docchatlink.constant.MethodEnum;
import com.docchatlink.dto.v2.ChunksResp;
import com.docchatlink.dto.v2.FileInfoResp;
import com.docchatlink.dto.v2.ResourceContentVO;
import com.docchatlink.iflytek2.client.IflytekDocV2Client;
import com.docchatlink.util.SpringUtils;
import com.docchatlink.vo.AnswerChatVO;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Slf4j
public class MyWebSocket2 extends WebSocketListener {

    // 断开websocket标志位
    private boolean wsCloseFlag;

    // 答案 将大模型返回结果全部接收，在组装成一句话返回
    private String totalAnswer = "";
    private String sid;
    private WebSocket webSocket;

    private WebSocketServer2 webSocketServer;


    public MyWebSocket2() {
        super();
    }

    public MyWebSocket2 create(Request request, OkHttpClient client, String sid, WebSocketServer2 webSocketServer) {
        this.wsCloseFlag = false;
        this.sid = sid;
        this.webSocket = client.newWebSocket(request, this);
        this.totalAnswer = "";
        this.webSocketServer = webSocketServer;
        return this;
    }

    @Override
    public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
        super.onOpen(webSocket, response);
        log.info("与讯飞星火知识库连接成功");
    }

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
        JSONObject result = JSONObject.parseObject(text);
        log.info("myJsonParse:{}", JSON.toJSONString(result));

        Integer status = result.getInteger("status");
        String content = result.getString("content");
        String xhSid = result.getString("sid");
        try {
            if (99 == status && null != result.get("sortedFileRefer")) {
                IflytekDocV2Client iflytekDocV2Client = SpringUtils.getBean(IflytekDocV2Client.class);
                List<ResourceContentVO> list = new ArrayList<>();
                List<Map<String, List<Integer>>> mapList = (List<Map<String, List<Integer>>>) result.get("sortedFileRefer");
                for (Map<String, List<Integer>> map : mapList) {
                    for (String key : map.keySet()) {
                        List<Integer> index = map.get(key);
                        ChunksResp chunksResp = iflytekDocV2Client.fileChunksWeb(key);
                        FileInfoResp fileInfoResp = iflytekDocV2Client.fileInfoWeb(key);
                        List<ChunksResp.DataVo> data = chunksResp.getData();

                        List<String> content1 = new ArrayList<>();
                        List<ResourceContentVO.IndexContentVO> indexContentVOS = new ArrayList<>();
                        for (ChunksResp.DataVo dataVo : data) {
                            if (index.contains(dataVo.getDataIndex())) {
                                byte[] utfBytes = new byte[0];
                                utfBytes = dataVo.getContent().getBytes(StandardCharsets.UTF_8);
                                String content2 = new String(utfBytes);
                                content1.add(content2);
                                ResourceContentVO.IndexContentVO indexContentVO = new ResourceContentVO.IndexContentVO();
                                indexContentVO.setIndex(dataVo.getDataIndex());
                                indexContentVO.setContent(content2);
                                indexContentVOS.add(indexContentVO);
                            }
                        }
                        ResourceContentVO resourceContentVO = new ResourceContentVO();
                        resourceContentVO.setFileId(key);
                        resourceContentVO.setFileName(fileInfoResp.getData().getFileName());
                        resourceContentVO.setContent(content1);
                        resourceContentVO.setIndexContents(indexContentVOS);
                        list.add(resourceContentVO);
                    }
                }
                log.info("来源返回 fileRefer：{}", result.get("sortedFileRefer"));
                AnswerChatVO answerChatVO = new AnswerChatVO()
                        .setMethodEnum(MethodEnum.RESOURCE)
                        .setContent(JSON.toJSONString(list))
                        .setXhSid(xhSid)
                        .setChatId(sid)
                        .setIsOver(true);
                webSocketServer.sendToOne(sid, JSONObject.toJSONString(answerChatVO));
            }
            if (2 == status) {
                byte[] utfBytes = new byte[0];
                utfBytes = content.getBytes(StandardCharsets.UTF_8);
                totalAnswer += new String(utfBytes);
                log.info("最后一次返回 totalAnswer：{}", totalAnswer);
                AnswerChatVO answerChatVO = new AnswerChatVO()
                        .setMethodEnum(MethodEnum.CALLED)
                        .setContent(content)
                        .setXhSid(xhSid)
                        .setChatId(sid)
                        .setIsOver(true);
                webSocketServer.sendToOne(sid, JSONObject.toJSONString(answerChatVO));

                // 处理所有结果的拼接
                AnswerChatVO allChatVO = new AnswerChatVO()
                        .setMethodEnum(MethodEnum.RESP_ALL)
                        .setContent(totalAnswer)
                        .setXhSid(xhSid)
                        .setChatId(sid)
                        .setIsOver(true);
                webSocketServer.sendToOne(sid, JSONObject.toJSONString(allChatVO));
            }
            if (1 == status) {
                log.info("第一次返回以及中间返回");
                byte[] utfBytes = new byte[0];
                utfBytes = content.getBytes(StandardCharsets.UTF_8);
                totalAnswer += new String(utfBytes);
                AnswerChatVO answerChatVO = new AnswerChatVO()
                        .setMethodEnum(MethodEnum.CALLED)
                        .setChatId(sid)
                        .setXhSid(xhSid)
                        .setContent(content)
                        .setIsOver(false);
                webSocketServer.sendToOne(sid, JSONObject.toJSONString(answerChatVO));
                log.info("第一次返回以及中间返回 temp{}", totalAnswer);
            }
        } catch (Exception e) {
            log.error(" webSocket 返回异常 ", e);
            AnswerChatVO answerChatVO = new AnswerChatVO()
                    .setMethodEnum(MethodEnum.CALLED)
                    .setContent("系统异常")
                    .setChatId(sid)
                    .setXhSid(xhSid)
                    .setIsOver(true);
            webSocketServer.sendToOne(sid, JSONObject.toJSONString(answerChatVO));
        }
    }

    @Override
    public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, Response response) {
        super.onFailure(webSocket, t, response);
        try {
            if (null != response) {
                int code = response.code();
                log.info("onFailure code:" + code);
                log.info("onFailure body:" + response.body().string());
                if (101 != code) {
                    log.info("connection failed 连接异常");
                }
            }
        } catch (IOException e) {
            log.error(" WebSocket connection failed ");
        }
    }

    @Override
    public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String mess) {
        try {
            log.info("连接关闭成功 ==> 关闭该连接信息");
        } catch (Exception e) {
            log.error(" WebSocket onClosed failed ", e);
            throw new RuntimeException(e);
        }
    }

    public void sendMessage(String mess) {
        log.info("发送服务器的消息为 {}", mess);
        totalAnswer = "";
        this.webSocket.send(mess);
    }

}
